// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <google/protobuf/arena.h>
#include <string>
#include <typeinfo>

#include <pollux/core/plan_node.h>
#include <pollux/type/type.h>

#include <pollux/substrait/substrait_extension_collector.h>
#include <pollux/substrait/pollux_to_substrait_expr.h>
#include <pollux/substrait/proto/substrait/algebra.pb.h>
#include <pollux/substrait/proto/substrait/plan.pb.h>

namespace kumo::pollux::substrait {
    /// Convert the Pollux plan into Substrait plan.
    class PolluxToSubstraitPlanConvertor {
    public:
        /// Convert Pollux PlanNode into Substrait Plan.
        /// @param vPlan Pollux query plan to convert.
        /// @param arena Arena to use for allocating Substrait plan objects.
        /// @return A pointer to Substrait plan object allocated on the arena and
        /// representing the input Pollux plan.
        ::substrait::Plan &toSubstrait(
            google::protobuf::Arena &arena,
            const core::PlanNodePtr &planNode);

    private:
        /// Convert Pollux PlanNode into Substrait Rel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const core::PlanNodePtr &planNode,
            ::substrait::Rel *rel);

        /// Convert Pollux FilterNode into Substrait FilterRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::FilterNode> &filterNode,
            ::substrait::FilterRel *filterRel);

        /// Convert Pollux ValuesNode into Substrait ReadRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::ValuesNode> &valuesNode,
            ::substrait::ReadRel *readRel);

        /// Convert Pollux ProjectNode into Substrait ProjectRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::ProjectNode> &projectNode,
            ::substrait::ProjectRel *projectRel);

        /// Convert Pollux Aggregation Node into Substrait AggregateRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::AggregationNode> &aggregateNode,
            ::substrait::AggregateRel *aggregateRel);

        /// Convert Pollux OrderBy Node into Substrait SortRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::OrderByNode> &orderByNode,
            ::substrait::SortRel *sortRel);

        /// Convert Pollux TopN Node into Substrait SortRel->FetchRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::TopNNode> &topNNode,
            ::substrait::FetchRel *fetchRel);

        /// Helper function to process sortingKeys and sortingOrders in Pollux to
        /// convert them to the sortField of SortRel in Substrait.
        const ::substrait::SortRel &processSortFields(
            google::protobuf::Arena &arena,
            const std::vector<core::FieldAccessTypedExprPtr> &sortingKeys,
            const std::vector<core::SortOrder> &sortingOrders,
            const RowTypePtr &inputType);

        /// Convert Pollux Limit Node into Substrait FetchRel.
        void toSubstrait(
            google::protobuf::Arena &arena,
            const std::shared_ptr<const core::LimitNode> &limitNode,
            ::substrait::FetchRel *fetchRel);

        /// Check there only have one source for the pollux node and return it.
        const core::PlanNodePtr &getSingleSource(const core::PlanNodePtr &node);

        /// The Expression converter used to convert Pollux representations into
        /// Substrait expressions.
        PolluxToSubstraitExprConvertorPtr exprConvertor_;

        /// The Type converter used to conver pollux representation into Substrait
        /// type.
        std::shared_ptr<PolluxToSubstraitTypeConvertor> typeConvertor_;

        /// The Extension collector storing the relations between the function
        /// signature and the function reference number.
        SubstraitExtensionCollectorPtr extensionCollector_;
    };
} // namespace kumo::pollux::substrait
